import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
from torch.optim.lr_scheduler import StepLR
from torchvision.utils import save_image
from torchvision.utils import make_grid
from torchmetrics.image.fid import FrechetInceptionDistance
import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image
import random
from tqdm import tqdm
from sklearn.metrics import confusion_matrix
from sklearn.manifold import TSNE
import seaborn as sns
import wandb

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device Used: {device}")

Device Used: cuda

saved_model_folder = "./saved_models"
os.makedirs(saved_model_folder, exist_ok=True)

class AFHQDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        """
        root_dir: path of the parent directory that contains images.
        transforms: augmentations applied to the images (can be none or more).
        """
        self.root_dir = root_dir
        self.transform = transform
        self.image_paths = []
        self.labels = []
        self.class_mapping = {}
        
        extensions = (".jpg", ".jpeg", ".png")
        # go through all sub-directories
        for label, category in enumerate(sorted(os.listdir(root_dir))):
            full_path = os.path.join(root_dir, category)
            if os.path.exists(full_path):
                self.class_mapping[label] = category
                for img_name in os.listdir(full_path):
                    if img_name.endswith(extensions):
                        self.image_paths.append(os.path.join(full_path, img_name))
                        self.labels.append(label)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        label = self.labels[idx]
        
        if self.transform:
            image = self.transform(image)
        
        return image, label

# Dataset Hyperparameters
img_size = 64
batch_size = 64

# dataset paths
train_dir = '/home/user/javeda1/stargan-v2/data/afhq/train'
val_dir = '/home/user/javeda1/stargan-v2/data/afhq/val'

transform = transforms.Compose([
    
        # transforms.GaussianBlur(3, sigma=(0.1, 2.0)),
        transforms.Resize((img_size, img_size)), # image is downsampled to 64x64
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

# Load the train and val dataset
train_dataset = AFHQDataset(root_dir=train_dir, transform=transform)
val_dataset = AFHQDataset(root_dir=val_dir, transform=transform)

# DataLoaders for train and val sets
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

train_size = len(train_loader.dataset)
val_size = len(val_loader.dataset)

print(f"Train dataset size: {train_size}")
print(f"Validation dataset size: {val_size}")

Train dataset size: 14630
Validation dataset size: 1500

def visualize_dataset(dataloader, num_images=12):
    """
    Visualizes a given number of images from the given data loader.
    """
    data_iter = iter(dataloader)
    images, labels = next(data_iter)
    class_mapping = dataloader.dataset.class_mapping
    
    # creates a suitable grid
    # grid_size = int(np.ceil(np.sqrt(num_images)))
    num_rows = int(np.sqrt(num_images))
    num_cols = int(np.ceil(num_images / num_rows))
    # print(num_cols, num_rows)
    plt.figure(figsize=(6, 6))
    plt.suptitle("Sample images from AFHQ dataset")
    
    for idx in range(min(num_images, len(images))):
        plt.subplot(num_rows, num_cols, idx + 1)
        
        img = images[idx].permute(1, 2, 0).numpy()
        max_val = np.max(img)
        min_val = np.min(img)
        img = (img - min_val) / (max_val - min_val)
        
        plt.imshow(img)
        plt.title(f'Class: {class_mapping[labels[idx].item()]}')
        plt.axis('off')
    
    plt.tight_layout()
    plt.show()

visualize_dataset(train_loader, num_images=20)

def vae_loss_function(recon_x, x, mu, log_var, kl_weight=1):
    """
    Calculates the VAE loss as a combination of 
    reconstruction loss and KL divergence, 
    scaled by a weight.
    """
    recon_loss = F.mse_loss(recon_x, x, reduction='sum')
    kl_divergence = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return recon_loss + kl_divergence * kl_weight, recon_loss, kl_divergence

def train_vae(model, train_loader, optimizer, kl_weight, device):
    """
    Trains the Variational Autoencoder (VAE) for one epoch on the given training data loader.
    """
    model.train()
    running_loss = 0.0
    running_recon_loss = 0.0
    running_kl_loss = 0.0
    
    for inputs, _ in train_loader:  #labels not used
        inputs = inputs.to(device)
        optimizer.zero_grad()
        recon_x, mu, log_var = model(inputs)
        loss, recon_loss, kl_loss = vae_loss_function(recon_x, inputs, mu, log_var, kl_weight)
        loss.backward()
        optimizer.step()
        
        # Collect all losses
        running_loss += loss.item()
        running_recon_loss += recon_loss.item()
        running_kl_loss += kl_loss.item()
        
    dataset_size = len(train_loader.dataset)
    return {
        'total_loss': running_loss / dataset_size,
        'recon_loss': running_recon_loss / dataset_size,
        'kl_loss': running_kl_loss / dataset_size
    }

def evaluate_vae(model, val_loader, kl_weight, device):
    """
    Evaluates the Variational Autoencoder (VAE) on the validation dataset after each epoch.
    """
    model.eval()
    running_loss = 0.0
    running_recon_loss = 0.0
    running_kl_loss = 0.0
    
    with torch.no_grad():
        for inputs, _ in val_loader:
            inputs = inputs.to(device)
            recon_x, mu, log_var = model(inputs)
            loss, recon_loss, kl_loss = vae_loss_function(recon_x, inputs, mu, log_var, kl_weight)
            
            # Collect all losses
            running_loss += loss.item()
            running_recon_loss += recon_loss.item()
            running_kl_loss += kl_loss.item()
            
    dataset_size = len(val_loader.dataset)
    return {
        'total_loss': running_loss / dataset_size,
        'recon_loss': running_recon_loss / dataset_size,
        'kl_loss': running_kl_loss / dataset_size
    }

def run_vae_training(
    model, train_loader, val_loader, device, num_epochs, learning_rate=0.001, 
    project="vae-training", name="vae_run", kl_weight=0.1, step_size=30, gamma=0.1):

    """Train and evaluate the model for a given number of epochs with W&B logging"""

    print(f"Training Name: {name}")
    print(f"Total num. of Epochs: {num_epochs}")
    print(f"Learning Rate: {learning_rate}")
    print(f"KL Weight used for Loss function: {kl_weight}\n")
    
    # Sample a batch for visualization ( used here to make it same for every epoch)
    inputs, _ = next(iter(train_loader))
    inputs = inputs.to(device)
    
    # Initialize W&B logging
    wandb.init(project=project, name=name, 
               config={
                   "learning_rate": learning_rate,
                   "num_epochs": num_epochs,
                   "step_size": step_size,
                   "gamma": gamma,
                   "kl_weight": kl_weight
               })
    
    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-5) # Optimizer
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)  # Learning rate scheduler
    
    for epoch in tqdm(range(num_epochs)):
        # Train for one epoch
        train_metrics = train_vae(model, train_loader, optimizer, kl_weight, device)
        # Evaluate after each epoch
        val_metrics = evaluate_vae(model, val_loader, kl_weight, device)
        
        current_lr = scheduler.get_last_lr()[0]
        scheduler.step()

        # Log images to W&B
        with torch.no_grad():
            # # Sample a small batch for visualization
            # inputs, _ = next(iter(train_loader))
            # inputs = inputs[:32].to(device)
            recon_x, _, _ = model(inputs)

            # Normalize and convert to image format
            recon_x = recon_x.view(-1, *inputs.shape[1:])
            recon_grid = make_grid(recon_x.cpu().detach() * 0.5 + 0.5, normalize=True, pad_value=1, padding=10)
            original_grid = make_grid(inputs.cpu().detach() * 0.5 + 0.5, normalize=True, pad_value=1, padding=10)
    
        # Log all data to W&B
        wandb.log({
            "epoch": epoch + 1,
            "train/total_loss": train_metrics['total_loss'],
            "train/recon_loss": train_metrics['recon_loss'],
            "train/kl_loss": train_metrics['kl_loss'],
            "val/total_loss": val_metrics['total_loss'],
            "val/recon_loss": val_metrics['recon_loss'],
            "val/kl_loss": val_metrics['kl_loss'],
            "learning_rate": current_lr,
            "original_images": wandb.Image(original_grid),
            "reconstructed_images": wandb.Image(recon_grid),
        })
        
        # # Print stats after each epoch
        print(f"\nEpoch {epoch+1}/{num_epochs}")
        print(f"Learning Rate: {current_lr:.6f}")
        print(f"Train - Total: {train_metrics['total_loss']:.4f}, "
              f"Recon: {train_metrics['recon_loss']:.4f}, "
              f"KL: {train_metrics['kl_loss']:.4f}")
        print(f"Eval  - Total: {val_metrics['total_loss']:.4f}, "
              f"Recon: {val_metrics['recon_loss']:.4f}, "
              f"KL: {val_metrics['kl_loss']:.4f}")
    
    # End W&B run
    wandb.finish()
    
    return model

def compute_fid_score(model, data_loader, device, max_samples=1000):
    """
    Computes the Fréchet Inception Distance (FID) between real and generated data.
    Taken from : https://lightning.ai/docs/torchmetrics/stable/image/frechet_inception_distance.html
    """
    fid = FrechetInceptionDistance(feature=2048).to(device)  # Use InceptionV3(default) features
    model.eval()
    
    total_samples = 0
    with torch.no_grad():
        for inputs, _ in data_loader:
            inputs = inputs.to(device)

            # Generate reconstructions
            recon_x, _, _ = model(inputs)

            # Convert images to uint8 and scale to [0, 255]
            inputs = (inputs * 255).byte()  
            recon_x = (recon_x * 255).byte()

            fid.update(inputs, real=True)
            fid.update(recon_x, real=False)

            total_samples += inputs.size(0)
    
    # Compute the FID score
    fid_score = fid.compute().item()
    return fid_score

def load_model(model_path):
    """
    Load a saved model
    """
    model = torch.load(model_path)
    return model

def visualize_reconstructions(model, images, device, num_images=8):
    """Visualizes original and reconstructed images from the VAE"""
    model.eval()
    
    images = images[:num_images].to(device)
    
    with torch.no_grad():
        recon_images, _, _ = model(images)
    
    # Move to CPU and convert to numpy
    images = images.cpu().permute(0, 2, 3, 1).numpy()
    recon_images = recon_images.cpu().permute(0, 2, 3, 1).numpy()
    
    # min-max normalization
    images = np.array([(img - np.min(img))/(np.max(img) - np.min(img)) for img in images])
    recon_images = np.array([(img - np.min(img))/(np.max(img) - np.min(img)) for img in recon_images])
    
    plt.figure(figsize=(num_images * 2, 4))
    for i in range(num_images):
        plt.subplot(2, num_images, i + 1)
        plt.imshow(images[i])
        plt.axis("off")
        plt.title("Original")
        
        plt.subplot(2, num_images, i + 1 + num_images)
        plt.imshow(recon_images[i])
        plt.axis("off")
        plt.title("Reconstructed")
        
    plt.suptitle("Original vs Reconstructed Images", fontsize=16)
    plt.tight_layout()
    plt.show()

class VAEEncoder(nn.Module):
    def __init__(self, latent_dim):
        super(VAEEncoder, self).__init__()
        
        # Initial convolution
        self.conv_initial = nn.Conv2d(3, 64, 3, stride=1, padding=1)
        
        # Downsampling blocks
        self.conv1 = nn.Conv2d(64, 128, 4, stride=2, padding=1)    
        self.conv2 = nn.Conv2d(128, 256, 4, stride=2, padding=1)   
        self.conv3 = nn.Conv2d(256, 512, 4, stride=2, padding=1)   
        self.conv4 = nn.Conv2d(512, 1024, 4, stride=2, padding=1)
        
        # Batch normalization
        self.bn1 = nn.BatchNorm2d(128)
        self.bn2 = nn.BatchNorm2d(256)
        self.bn3 = nn.BatchNorm2d(512)
        self.bn4 = nn.BatchNorm2d(1024)
        
        
        # Bottleneck
        self.fc_mu = nn.Linear(1024 * 4 * 4, latent_dim)
        self.fc_var = nn.Linear(1024 * 4 * 4, latent_dim)
        
        # Dropout for regularization
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, x):
        # # Ensure input is 4D: [batch_size, channels, height, width]
        # if x.dim() == 3:
        #     x = x.unsqueeze(0)  # Add batch dimension if missing
        # elif x.dim() != 4:
        #     raise ValueError(f"Expected 3D or 4D input (got {x.dim()}D input)")
            
        # # Check channels dimension
        # if x.size(1) != 3:
        #     raise ValueError(f"Expected 3 channels (got {x.size(1)} channels)")
            
        # Initial convolution
        x = F.leaky_relu(self.conv_initial(x), 0.2)
        
        # Downsampling path
        x = F.leaky_relu(self.bn1(self.conv1(x)), 0.2)
        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2)
        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)
        x = F.leaky_relu(self.bn4(self.conv4(x)), 0.2)
        
        
        # Flatten and apply dropout
        x = self.dropout(x.view(x.size(0), -1))
        
        # Generate latent parameters
        mu = self.fc_mu(x)
        log_var = self.fc_var(x)
        
        return mu, log_var

class VAEDecoder(nn.Module):
    def __init__(self, latent_dim):
        super(VAEDecoder, self).__init__()
        
        # Initial fully connected layer
        self.fc = nn.Linear(latent_dim, 1024 * 4 * 4)
        
        # Upsampling blocks
        self.conv1 = nn.ConvTranspose2d(1024, 512, 4, stride=2, padding=1) 
        self.conv2 = nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1)  
        self.conv3 = nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1)  
        self.conv4 = nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1)   
        
        # Batch normalization
        # self.bna = nn.BatchNorm2d(1024)
        
        self.bn1 = nn.BatchNorm2d(512)
        self.bn2 = nn.BatchNorm2d(256)
        self.bn3 = nn.BatchNorm2d(128)
        self.bn4 = nn.BatchNorm2d(64)
        
        self.conv_final = nn.Conv2d(64, 3, 3, stride=1, padding=1)
        
        # Dropout
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, x):
        # Reshape from latent space
        x = F.relu(self.fc(x))
        x = x.view(x.size(0), 1024, 4, 4)
        
        # Upsampling path
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = F.relu(self.bn4(self.conv4(x)))
        
        # Final convolution with tanh activation
        x = torch.tanh(self.conv_final(x))
        
        return x

class ConvVAE(nn.Module):
    def __init__(self, latent_dim):
        super(ConvVAE, self).__init__()
        self.encoder = VAEEncoder(latent_dim)
        self.decoder = VAEDecoder(latent_dim)
        
    def reparameterize(self, mu, log_var):
        if self.training:
            std = torch.exp(0.5 * log_var)
            eps = torch.randn_like(std)
            return mu + eps * std
        return mu
        
    def forward(self, x):
        mu, log_var = self.encoder(x)
        z = self.reparameterize(mu, log_var)
        return self.decoder(z), mu, log_var

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=1

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: ")

Training Name: run_kl_wgt_1_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 1

wandb: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.
wandb: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)
wandb: You can find your API key in your browser here: https://wandb.ai/authorize
wandb: Paste an API key from your profile and hit enter, or press ctrl+c to quit:

wandb: Appending key for api.wandb.ai to your netrc file: /home/user/javeda1/.netrc

  2%|█▍                                                      | 1/40 [00:10<06:36, 10.17s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 1653.9196, Recon: 1500.0610, KL: 153.8587
Eval  - Total: 1645.1742, Recon: 1485.3813, KL: 159.7929

  5%|██▊                                                     | 2/40 [00:19<06:02,  9.53s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 1093.1337, Recon: 921.2636, KL: 171.8701
Eval  - Total: 1065.1776, Recon: 879.4257, KL: 185.7518

  8%|████▏                                                   | 3/40 [00:28<05:46,  9.38s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 968.1743, Recon: 780.8183, KL: 187.3559
Eval  - Total: 887.5771, Recon: 703.6158, KL: 183.9613

 10%|█████▌                                                  | 4/40 [00:37<05:35,  9.32s/it]

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)
Downloading: "https://github.com/toshas/torch-fidelity/releases/download/v0.2.0/weights-inception-2015-12-05-6726825d.pth" to /home/user/javeda1/.cache/torch/hub/checkpoints/weights-inception-2015-12-05-6726825d.pth
100%|███████████████████████████████████████████████████| 91.2M/91.2M [00:00<00:00, 114MB/s]

FID Score for model run_kl_wgt_1_ep_40_latent_dim_512: 147.1151885986328

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=0.1

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: {save_path}")

Training Name: run_kl_wgt_0.1_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 0.1

wandb: Currently logged in as: ahmadjaved97. Use `wandb login --relogin` to force relogin

  2%|█▍                                                      | 1/40 [00:10<06:30, 10.00s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 1288.4435, Recon: 1236.6438, KL: 517.9968
Eval  - Total: 832.0195, Recon: 777.5456, KL: 544.7392

  5%|██▊                                                     | 2/40 [00:20<06:27, 10.20s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 695.2330, Recon: 632.6024, KL: 626.3055
Eval  - Total: 591.3673, Recon: 526.3305, KL: 650.3685

  8%|████▏                                                   | 3/40 [00:30<06:21, 10.32s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 597.5521, Recon: 528.9844, KL: 685.6768
Eval  - Total: 535.0150, Recon: 466.5609, KL: 684.5407

 10%|█████▌                                                  | 4/40 [00:40<06:06, 10.18s/it]

Epoch 4/40
Learning Rate: 0.000100
Train - Total: 548.9701, Recon: 477.7421, KL: 712.2808
Eval  - Total: 503.0447, Recon: 432.1050, KL: 709.3967

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

FID Score for model run_kl_wgt_0.1_ep_40_latent_dim_512: 122.6679916381836

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=0.01

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: {save_path}")

Training Name: run_kl_wgt_0.01_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 0.01

  2%|█▍                                                      | 1/40 [00:10<07:03, 10.85s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 1181.5184, Recon: 1166.5779, KL: 1494.0456
Eval  - Total: 712.0218, Recon: 697.3104, KL: 1471.1411

  5%|██▊                                                     | 2/40 [00:20<06:32, 10.32s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 604.7859, Recon: 590.2766, KL: 1450.9332
Eval  - Total: 551.4014, Recon: 536.2569, KL: 1514.4445

  8%|████▏                                                   | 3/40 [00:31<06:32, 10.60s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 498.4249, Recon: 482.8365, KL: 1558.8424
Eval  - Total: 470.2254, Recon: 454.3379, KL: 1588.7437

 10%|█████▌                                                  | 4/40 [00:45<07:07, 11.87s/it]

Epoch 4/40
Learning Rate: 0.000100
Train - Total: 446.1589, Recon: 429.8908, KL: 1626.8100
Eval  - Total: 439.4034, Recon: 422.6276, KL: 1677.5747

 12%|███████                                                 | 5/40 [00:58<07:08, 12.25s/it]

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

FID Score for model run_kl_wgt_0.01_ep_40_latent_dim_512: 116.9962158203125

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=0.001

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: {save_path}")

Training Name: run_kl_wgt_0.001_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 0.001

  2%|█▍                                                      | 1/40 [00:09<06:01,  9.27s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 1205.5347, Recon: 1202.0513, KL: 3483.4063
Eval  - Total: 688.4095, Recon: 686.2839, KL: 2125.5886

  5%|██▊                                                     | 2/40 [00:18<05:52,  9.27s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 587.9639, Recon: 585.6678, KL: 2296.1871
Eval  - Total: 514.4709, Recon: 511.9555, KL: 2515.3943

  8%|████▏                                                   | 3/40 [00:27<05:42,  9.26s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 482.1384, Recon: 479.5534, KL: 2584.9820
Eval  - Total: 459.3856, Recon: 456.5578, KL: 2827.7235

 10%|█████▌                                                  | 4/40 [00:37<05:33,  9.27s/it]

Epoch 4/40
Learning Rate: 0.000100
Train - Total: 434.3168, Recon: 431.5614, KL: 2755.3510
Eval  - Total: 433.4627, Recon: 430.6415, KL: 2821.2776

 12%|███████                                                 | 5/40 [00:46<05:23,  9.24s/it]

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

FID Score for model run_kl_wgt_0.001_ep_40_latent_dim_512: 116.66500091552734

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=1.5

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: {save_path}")

Training Name: run_kl_wgt_1.5_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 1.5

  2%|█▍                                                      | 1/40 [00:09<05:52,  9.05s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 1577.3921, Recon: 1408.9965, KL: 112.2637
Eval  - Total: 1548.4882, Recon: 1356.5669, KL: 127.9475

  5%|██▊                                                     | 2/40 [00:18<05:44,  9.07s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 1162.5180, Recon: 965.8803, KL: 131.0918
Eval  - Total: 1194.2998, Recon: 988.1949, KL: 137.4033

  8%|████▏                                                   | 3/40 [00:27<05:35,  9.07s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 1052.7086, Recon: 842.1593, KL: 140.3662
Eval  - Total: 1187.0521, Recon: 986.1494, KL: 133.9352

 10%|█████▌                                                  | 4/40 [00:36<05:27,  9.10s/it]

Epoch 4/40
Learning Rate: 0.000100
Train - Total: 994.5136, Recon: 783.0087, KL: 141.0032
Eval  - Total: 984.5990, Recon: 775.4794, KL: 139.4131

 12%|███████                                                 | 5/40 [00:45<05:18,  9.11s/it]

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

FID Score for model run_kl_wgt_1.5_ep_40_latent_dim_512: 157.549072265625

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=2

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: {save_path}")

Training Name: run_kl_wgt_2_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 2

  2%|█▍                                                      | 1/40 [00:09<05:59,  9.23s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 1680.3621, Recon: 1487.6368, KL: 96.3627
Eval  - Total: 2834.4174, Recon: 2626.1163, KL: 104.1506

  5%|██▊                                                     | 2/40 [00:18<05:51,  9.24s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 1249.2274, Recon: 1034.5994, KL: 107.3140
Eval  - Total: 1530.9900, Recon: 1313.6945, KL: 108.6478

  8%|████▏                                                   | 3/40 [00:27<05:42,  9.27s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 1138.4374, Recon: 908.4061, KL: 115.0157
Eval  - Total: 1128.0360, Recon: 892.3010, KL: 117.8675

 10%|█████▌                                                  | 4/40 [00:37<05:33,  9.27s/it]

Epoch 4/40
Learning Rate: 0.000100
Train - Total: 1073.9766, Recon: 841.4963, KL: 116.2402
Eval  - Total: 1014.8824, Recon: 779.6643, KL: 117.6090

 12%|███████                                                 | 5/40 [00:46<05:23,  9.24s/it]

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

FID Score for model run_kl_wgt_2_ep_40_latent_dim_512: 162.39413452148438

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=5

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: {save_path}")

Training Name: run_kl_wgt_5_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 5

Traceback (most recent call last):
  File "/home/user/javeda1/miniconda3/lib/python3.12/multiprocessing/util.py", line 303, in _run_finalizers
    finalizer()
  File "/home/user/javeda1/miniconda3/lib/python3.12/multiprocessing/util.py", line 227, in __call__
    res = self._callback(*self._args, **self._kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/javeda1/miniconda3/lib/python3.12/multiprocessing/util.py", line 136, in _remove_temp_dir
    rmtree(tempdir, onerror=onerror)
  File "/home/user/javeda1/miniconda3/lib/python3.12/shutil.py", line 759, in rmtree
    _rmtree_safe_fd(stack, onexc)
  File "/home/user/javeda1/miniconda3/lib/python3.12/shutil.py", line 703, in _rmtree_safe_fd
    onexc(func, path, err)
  File "/home/user/javeda1/miniconda3/lib/python3.12/shutil.py", line 750, in onexc
    return onerror(func, path, exc_info)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/javeda1/miniconda3/lib/python3.12/shutil.py", line 662, in _rmtree_safe_fd
    os.rmdir(name, dir_fd=dirfd)
OSError: [Errno 39] Directory not empty: '/tmp/pymp-o5yehjy3'

  2%|█▍                                                      | 1/40 [00:09<06:00,  9.23s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 2012.3246, Recon: 1740.1891, KL: 54.4271
Eval  - Total: 5183.5230, Recon: 4913.3266, KL: 54.0393

  5%|██▊                                                     | 2/40 [00:18<05:51,  9.24s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 1529.6669, Recon: 1254.5351, KL: 55.0264
Eval  - Total: 3294.9925, Recon: 3011.1629, KL: 56.7659

  8%|████▏                                                   | 3/40 [00:27<05:43,  9.29s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 1442.3994, Recon: 1156.7717, KL: 57.1255
Eval  - Total: 2145.4898, Recon: 1843.2327, KL: 60.4514

 10%|█████▌                                                  | 4/40 [00:36<05:32,  9.24s/it]

Epoch 4/40
Learning Rate: 0.000100
Train - Total: 1356.8731, Recon: 1069.6204, KL: 57.4505
Eval  - Total: 1902.2592, Recon: 1620.5047, KL: 56.3509

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

FID Score for model run_kl_wgt_5_ep_40_latent_dim_512: 179.6144561767578

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

# Model training parameters
learning_rate=0.0001
step_size=10
gamma=0.5

kl_weight=10

num_epochs=40

latent_dim = 512 # define latent dimension

name=f"run_kl_wgt_{str(kl_weight)}_ep_{num_epochs}_latent_dim_{latent_dim}"
project="assignment-5-v2"


# Load model
model = ConvVAE(latent_dim=latent_dim).to(device)
# model

model = run_vae_training(
    model, train_loader, val_loader, device, 
    num_epochs=num_epochs, learning_rate=learning_rate,
    step_size=step_size, gamma=gamma,
    kl_weight=kl_weight,
    name=name, project=project
)

save_path = os.path.join(saved_model_folder, name)
torch.save(model, save_path)
print(f"Model saved at: {save_path}")

Training Name: run_kl_wgt_10_ep_40_latent_dim_512
Total num. of Epochs: 40
Learning Rate: 0.0001
KL Weight used for Loss function: 10

  2%|█▍                                                      | 1/40 [00:09<05:55,  9.13s/it]

Epoch 1/40
Learning Rate: 0.000100
Train - Total: 2157.0848, Recon: 1825.2464, KL: 33.1838
Eval  - Total: 6552.8518, Recon: 6189.9753, KL: 36.2876

  5%|██▊                                                     | 2/40 [00:18<05:52,  9.27s/it]

Epoch 2/40
Learning Rate: 0.000100
Train - Total: 1755.9227, Recon: 1419.9669, KL: 33.5956
Eval  - Total: 5520.5934, Recon: 5141.5175, KL: 37.9076

  8%|████▏                                                   | 3/40 [00:27<05:43,  9.28s/it]

Epoch 3/40
Learning Rate: 0.000100
Train - Total: 1693.6546, Recon: 1358.3253, KL: 33.5329
Eval  - Total: 4729.6289, Recon: 4400.4534, KL: 32.9175

 10%|█████▌                                                  | 4/40 [00:36<05:31,  9.21s/it]

Epoch 4/40
Learning Rate: 0.000100
Train - Total: 1644.8664, Recon: 1309.4471, KL: 33.5419
Eval  - Total: 3248.0873, Recon: 2969.4295, KL: 27.8658

 12%|███████                                                 | 5/40 [00:46<05:24,  9.26s/it]

model = load_model(save_path)
max_samples = min(2000, len(val_loader.dataset))
fid_score = compute_fid_score(model, val_loader, device, max_samples=max_samples)
print(f"FID Score for model {name}: {fid_score}")

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

FID Score for model run_kl_wgt_10_ep_40_latent_dim_512: 186.1881561279297

data_iter = iter(val_loader)
images, _ = next(data_iter)
visualize_reconstructions(model, images, device)

def visualize_latent_space_tsne(model, dataloader, num_samples=2000, device="cuda" if torch.cuda.is_available() else "cpu"):
    """
    Visualize latent space by encoding images and projecting latent representations using T-SNE.
    """
    model.eval()
    latents = []
    labels = []

    with torch.no_grad():
        for i, (img_batch, label_batch) in enumerate(dataloader):
            if len(latents) >= num_samples:
                break
            img_batch = img_batch.to(device)
            mu, _ = model.encoder(img_batch)
            latents.append(mu.cpu().numpy())
            labels.extend(label_batch.cpu().numpy())
    
    # Flatten for TSNE visualization
    latents = np.concatenate(latents)[:num_samples]
    labels = np.array(labels)[:num_samples]

    # Apply T-SNE
    tsne_latents = TSNE(n_components=2).fit_transform(latents)

    # Plot the projection
    plt.figure(figsize=(8, 6))
    scatter = plt.scatter(tsne_latents[:, 0], tsne_latents[:, 1], c=labels, cmap='viridis', alpha=0.5)
    plt.title("T-SNE Projection of Latent Representations")
    plt.xlabel("T-SNE Component 1")
    plt.ylabel("T-SNE Component 2")
    
    plt.legend(handles=scatter.legend_elements()[0], labels=['cat', 'dog', 'wild'], title="Classes")
    plt.show()

# model used here is with kl_weight=10
visualize_latent_space_tsne(model, val_loader)

def visualize_image_tsne(model, dataloader, num_samples=2000, device="cuda" if torch.cuda.is_available() else "cpu"):
    """
    Visualize images by projecting image pixel data using T-SNE.
    """
    model.eval()
    images = []
    labels = []

    with torch.no_grad():
        for i, (img_batch, label_batch) in enumerate(dataloader):
            if len(images) >= num_samples:
                break
            images.append(img_batch.numpy())
            labels.extend(label_batch.numpy())

    images = np.concatenate(images)[:num_samples]
    labels = np.array(labels)[:num_samples]

    # Flatten for TSNE visualization
    images_flat = images.reshape(images.shape[0], -1)

    # Apply T-SNE
    tsne_imgs = TSNE(n_components=2).fit_transform(images_flat)

    # Plot the projection
    plt.figure(figsize=(8, 6))
    scatter = plt.scatter(tsne_imgs[:, 0], tsne_imgs[:, 1], c=labels, cmap='viridis', alpha=0.5)
    plt.title("T-SNE Projection of Images")
    plt.xlabel("T-SNE Component 1")
    plt.ylabel("T-SNE Component 2")
    
    plt.legend(handles=scatter.legend_elements()[0], labels=['cat', 'dog', 'wild'], title="Classes")
    
    plt.show()

# model used here is with kl_weight=10
visualize_image_tsne(model, val_loader)

def generate_images(model, latent_dim, num_images, device):
    """
    Generate new images by sampling latent vectors and decoding them with the VAE model.
    """
    model.eval()
    with torch.no_grad():
        # Sample random latent vectors
        z = torch.randn(num_images, latent_dim).to(device)
        generated_images = model.decoder(z)
    return generated_images

def visualize_latent_interpolation_grid(model, latent_dim, num_steps, device):
    """
    Visualize interpolations in the latent space between four random points as a grid.
    """
    model.eval()
    with torch.no_grad():
        # Sample four random latent vectors
        z1 = torch.randn(1, latent_dim).to(device)
        z2 = torch.randn(1, latent_dim).to(device)
        z3 = torch.randn(1, latent_dim).to(device)
        z4 = torch.randn(1, latent_dim).to(device)

        # Interpolation weights
        weights = torch.linspace(0, 1, num_steps).to(device)

        
        grid_images = []
        for alpha in weights:
            row = []
            for beta in weights:
                # Bilinear interpolation
                z = (1 - alpha) * (1 - beta) * z1 + \
                    (1 - alpha) * beta * z2 + \
                    alpha * (1 - beta) * z3 + \
                    alpha * beta * z4
                row.append(model.decoder(z))
            grid_images.append(torch.cat(row, dim=0))

        # Concatenate rows to form the grid
        grid_images = torch.cat(grid_images, dim=0)
        

        grid_images = (grid_images * 0.5 + 0.5).cpu()
        grid = make_grid(grid_images, nrow=num_steps, normalize=True, pad_value=1)

        # Plot the grid
        plt.figure(figsize=(15, 15))
        plt.imshow(grid.permute(1, 2, 0))
        plt.axis('off')
        plt.title("Latent Space Interpolation")
        plt.show()

num_images = 16
latent_dim = 512  # Adjust based on your model
generated_images = generate_images(model, latent_dim, num_images, device).to("cpu")

# Visualize generated images
grid = make_grid(generated_images * 0.5 + 0.5, nrow=5, normalize=True, pad_value=1, padding=5)
plt.figure(figsize=(10, 5))
plt.imshow(grid.permute(1, 2, 0))
plt.axis('off')
plt.title("Generated Images")
plt.show()

num_steps = 10  # Number of interpolation steps
visualize_latent_interpolation_grid(model, latent_dim, num_steps, device)

num_steps = 10  # Number of interpolation steps
model = load_model("./saved_models/run_kl_wgt_0.1_ep_40_latent_dim_512")
visualize_latent_interpolation_grid(model, latent_dim, num_steps, device)

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

model = load_model("./saved_models/run_kl_wgt_0.1_ep_40_latent_dim_512")
num_images = 16
latent_dim = 512  # Adjust based on your model
generated_images = generate_images(model, latent_dim, num_images, device).to("cpu")

# Visualize generated images
grid = make_grid(generated_images * 0.5 + 0.5, nrow=5, normalize=True, pad_value=1, padding=5)
plt.figure(figsize=(10, 5))
plt.imshow(grid.permute(1, 2, 0))
plt.axis('off')
plt.title("Generated Images")
plt.show()

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

def interpolate_latent(model, dataloader, num_steps=8, device="cuda"):
    """
    Interpolate between two random images in the latent space and visualize the interpolation.
    """
    model.eval()
    
    # Randomly pick two images
    batch_size = 2 
    idx = torch.randint(0, len(dataloader.dataset), (batch_size,))  # Random indices
    images = [dataloader.dataset[i][0] for i in idx]  # Fetch images 
    img1, img2 = images[0].to(device), images[1].to(device) 
    
    with torch.no_grad():
        # Encode the two randomly chosen images
        mu1, log_var1 = model.encoder(img1.unsqueeze(0))
        mu2, log_var2 = model.encoder(img2.unsqueeze(0))
        
        # Generate interpolation steps
        alphas = np.linspace(0, 1, num_steps)
        interp_images = []
        
        for alpha in alphas:
            # Perform linear interpolation in latent space
            z = mu1 * (1 - alpha) + mu2 * alpha
            decoded = model.decoder(z)
            interp_images.append(decoded)
        
        # Concatenate the interpolated images
        interp_images = torch.cat(interp_images).cpu().permute(0, 2, 3, 1).numpy()
        interp_images = (interp_images + 1) / 2 
        
        # Plot the interpolated images
        plt.figure(figsize=(2 * num_steps, 2))
        for i in range(num_steps):
            plt.subplot(1, num_steps, i + 1)
            plt.imshow(np.clip(interp_images[i], 0, 1))
            plt.axis('off')
        plt.show()

interpolate_latent(model, val_loader, num_steps=10)

model = load_model("./saved_models/run_kl_wgt_0.1_ep_40_latent_dim_512")
interpolate_latent(model, val_loader, num_steps=10)

/tmp/ipykernel_1886659/2494337095.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  model = torch.load(model_path)

epoch	▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
learning_rate	██████████▄▄▄▄▄▄▄▄▄▄▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁
train/kl_loss	▁▅███▇▇▆▆▅▅▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄
train/recon_loss	█▄▃▃▃▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
train/total_loss	█▄▃▃▃▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
val/kl_loss	▁▇▇█▆▅▆▆▄▄▂▄▂▃▃▂▃▁▂▃▂▂▂▂▃▂▃▂▃▂▂▂▂▃▃▂▃▂▂▂
val/recon_loss	█▄▃▂▃▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
val/total_loss	█▄▃▃▃▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁

epoch	40
learning_rate	1e-05
train/kl_loss	170.14351
train/recon_loss	437.61732
train/total_loss	607.76083
val/kl_loss	164.7427
val/recon_loss	398.64764
val/total_loss	563.39034

Model	FID Score
run_kl_wgt_1_ep_40_latent_dim_512	147.12
run_kl_wgt_0.1_ep_40_latent_dim_512	122.67
run_kl_wgt_0.01_ep_40_latent_dim_512	116.99
run_kl_wgt_0.001_ep_40_latent_dim_512	116.67
run_kl_wgt_1.5_ep_40_latent_dim_512	157.55
run_kl_wgt_2_ep_40_latent_dim_512	162.39
run_kl_wgt_5_ep_40_latent_dim_512	179.61
run_kl_wgt_10_ep_40_latent_dim_512	186.19

Inspired from¶

Training logs¶

Import required libraries¶

Set Device¶

Data Loading and Processing¶

Define augmentations¶

Visualize Dataset¶

Training and Validation Code¶

FID Score¶

Load saved model¶

Result visualizations¶

Model Architecture¶

Model training¶

1. Model with kl_weight=1¶

Run history:

Run summary:

2. Model with kl_weight=0.1¶

Run history:

Run summary:

3. Model with kl_weight=0.01¶

Run history:

Run summary:

4. Model with kl_weight=0.001¶

Run history:

Run summary:

5. Model with kl_weight=1.5¶

Run history:

Run summary:

6. Model with kl_weight=2¶

Run history:

Run summary:

7. Model with kl_weight=5¶

Run history:

Run summary:

7. Model with kl_weight=10¶

Run history:

Run summary:

Latent Space and Image TSNE Projection¶

Generating new images from latent dimension¶

Latent Space Interpolation¶

For model with kl_weight=10¶

For model with kl_weight=10¶

For model with kl_weight=0.1¶

Interpolation between 2 images¶

Analysis of the trained models¶

FID scores of models¶

Qualitative Comparison¶

Loss Analysis¶